🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20226

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ROCm] MI350 CI jobs will have longer queue times due to CI migration

❌ 1 New Failure, 2 Pending, 2 Unrelated Failures

As of commit 8e81436 with merge base 5526971 ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

@claude review

Claude finished @JulianCloudNTH's task in 4m 47s —— View job

Claude Code Review

• Read changed files
• Review WGSL kernel
• Review C++ dispatch & resource management
• Cross-check against backend conventions (SDPA, codegen)
• Post feedback

Overall this is a clean, well-guarded first cut. Validation-before-allocation, the fp32-only byte-size checks, the K_packed == ceil(K/2) guard, and the resource-release pattern all match the existing SDPA op conventions, and the generated header's wgsl-sha256 matches the .wgsl source (verified locally), so codegen is in sync. Comments are appropriately minimal. A few correctness edge cases are worth tightening before this lands.

Potential bugs / correctness

1. Weight bound as array<u32> can over-read the last word when N * K_packed isn't a multiple of 4 — QuantizedLinear.cpp:198, q4gsw_linear.wgsl:45

The weight is bound as array<u32> with size = weight.nbytes = N * K_packed bytes. The shader computes t_weight[byte_idx >> 2u] with a maximum byte_idx = N*K_packed - 1, so the highest word index is (N*K_packed - 1) >> 2. When N * K_packed is not a multiple of 4, the effective array<u32> length is floor(nbytes / 4), and that final word read falls outside the bound — WebGPU robustness clamps it to 0, silently producing wrong values for the affected output column(s). This is shape-dependent (e.g. any N*K_packed ≡ 1,2,3 (mod 4)), so it can pass on aligned test shapes and fail in the wild.

Consider validating (N * K_packed) % 4 == 0, or rounding the weight binding size up to a 4-byte multiple, or indexing the weight as array<u32> only after asserting alignment. Worth a dedicated test shape with an odd N*K_packed.

Fix this →

2. scales bounds (num_groups, padded_N) are not validated against the indices the shader uses — QuantizedLinear.cpp:69-70, q4gsw_linear.wgsl:54

The shader indexes t_scales[(k / group_size) * padded_N + n], with max row (K-1)/group_size and max col N-1. The current byte-size guard only checks scales.nbytes == num_groups * padded_N * 4; it never verifies that num_groups >= ceil(K / group_size) or that padded_N >= N. If either is smaller than expected (mismatched serialization), the kernel reads OOB scales (robustness-clamped, wrong results) rather than failing loudly. Adding if (num_groups < (K + group_size - 1) / group_size || padded_N < N) throw … would catch malformed inputs at dispatch time. Note group_size is read at line 96-102, so this check belongs after that.

3. M = in_numel / K truncates silently if in_numel % K != 0 — QuantizedLinear.cpp:66

A non-multiple input numel would silently drop the remainder into a smaller M rather than erroring. A quick if (in_numel % K != 0) throw … keeps the failure mode explicit, consistent with the other malformed-dims guards.

Performance (non-blocking)

4. Scale reload + integer division every K iteration — q4gsw_linear.wgsl:54

(k / params.group_size) * padded_N + n recomputes an integer division and reloads the scale on every inner iteration, even though the scale only changes when k crosses a group boundary. Since this is the matmul hot loop, hoisting the scale fetch to update only at group boundaries (or iterating group-by-group) would cut a division + global load per MAC. Fine to defer if this kernel is correctness-first for now, but worth a // TODO.

Minor / style

• QuantizedLinear.cpp:151 declares WGPUBindGroupLayoutEntry entries[6] and the SDPA op factored this into a shared build_dispatch helper (ops/sdpa/Sdpa.cpp:150). Not required, but if more quantized-linear variants are coming in the stack, reusing that helper would avoid duplicating the bind-group/pipeline boilerplate.
• static_assert(sizeof(Q4gswParams) == 32, …) is a nice guard against accidental struct drift — good.
• Error strings consistently prefixed WebGPU linear_q4gsw: — good for triage.

None of the above are blockers except items 1–2, which I'd recommend guarding (or covering with an odd-shape test) before relying on this in production, since both fail silently via robustness clamping rather than throwing.